import torch
from ultralytics.nn.modules.conv import Conv
from ultralytics.nn.modules.block import C2f # C2PSA, C3k2
from ultralytics.nn.extra_modules.prune_module import *
from ultralytics.nn.modules.MSAD import C2f_DynamicAxialRouterV2, C2f_DynamicAxialRouterV2_v2

# 没有购买yolov8项目需要注释以下
# from ultralytics.nn.extra_modules.block import C2f_Faster, C2f_EMBC, RepNCSPELAN4, C2f_Star

def transfer_weights_c2f_v2_to_c2f(c2f_v2, c2f):
    c2f.cv2 = c2f_v2.cv2
    c2f.m = c2f_v2.m

    state_dict = c2f.state_dict()
    state_dict_v2 = c2f_v2.state_dict()

    # Transfer cv1 weights from C2f to cv0 and cv1 in C2f_v2
    old_weight = state_dict['cv1.conv.weight']
    new_cv1 = Conv(c1=state_dict_v2['cv0.conv.weight'].size()[1],
                   c2=(state_dict_v2['cv0.conv.weight'].size()[0] + state_dict_v2['cv1.conv.weight'].size()[0]),
                   k=c2f_v2.cv1.conv.kernel_size,
                   s=c2f_v2.cv1.conv.stride)
    c2f.cv1 = new_cv1
    c2f.c1, c2f.c2 = state_dict_v2['cv0.conv.weight'].size()[0], state_dict_v2['cv1.conv.weight'].size()[0]
    state_dict['cv1.conv.weight'] = torch.cat([state_dict_v2['cv0.conv.weight'], state_dict_v2['cv1.conv.weight']], dim=0)

    # Transfer cv1 batchnorm weights and buffers from C2f to cv0 and cv1 in C2f_v2
    for bn_key in ['weight', 'bias', 'running_mean', 'running_var']:
        state_dict[f'cv1.bn.{bn_key}'] = torch.cat([state_dict_v2[f'cv0.bn.{bn_key}'], state_dict_v2[f'cv1.bn.{bn_key}']], dim=0)

    # Transfer remaining weights and buffers
    for key in state_dict:
        if not key.startswith('cv1.'):
            state_dict[key] = state_dict_v2[key]

    # Transfer all non-method attributes
    for attr_name in dir(c2f_v2):
        attr_value = getattr(c2f_v2, attr_name)
        if not callable(attr_value) and '_' not in attr_name:
            setattr(c2f, attr_name, attr_value)

    c2f.load_state_dict(state_dict)

def replace_c2f_v2_with_c2f(module):
    # for yolov8
    for name, child_module in module.named_children():
        if isinstance(child_module, C2f_v2):
            # Replace C2f with C2f_v2 while preserving its parameters
            shortcut = infer_shortcut(child_module.m[0])
            c2f = C2f_infer(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
                            n=len(child_module.m), shortcut=shortcut,
                            g=child_module.m[0].cv2.conv.groups,
                            e=child_module.c / child_module.cv2.conv.out_channels)
            transfer_weights_c2f_v2_to_c2f(child_module, c2f)
            setattr(module, name, c2f)
        else:
            replace_c2f_v2_with_c2f(child_module)
    
    # for name, child_module in module.named_children():
    #     if isinstance(child_module, C2f_EMBC_v2):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f = C2f_EMBC_infer(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=1,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_v2_to_c2f(child_module, c2f)
    #         setattr(module, name, c2f)
    #     elif isinstance(child_module, C2f_v2):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f = C2f_infer(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=child_module.m[0].cv2.conv.groups,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_v2_to_c2f(child_module, c2f)
    #         setattr(module, name, c2f)
    #     else:
    #         replace_c2f_v2_with_c2f(child_module)
    
    # for yolov10
    # for name, child_module in module.named_children():
    #     if isinstance(child_module, C2fCIB_v2):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f = C2fCIB_infer(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         lk=(False if type(child_module.m[0].cv1[2]) is Conv else True),
    #                         g=1,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_v2_to_c2f(child_module, c2f)
    #         setattr(module, name, c2f)
    #     elif isinstance(child_module, C2f_v2):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f = C2f_infer(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=child_module.m[0].cv2.conv.groups,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_v2_to_c2f(child_module, c2f)
    #         setattr(module, name, c2f)
    #     else:
    #         replace_c2f_v2_with_c2f(child_module)
    
    # for yolo11
    # for name, child_module in module.named_children():
    #     if isinstance(child_module, C3k2_v2):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c3k = child_module.m[0] is C3k
    #         if c3k:
    #             g = child_module.m[0].m[0].cv2.conv.groups
    #         else:
    #             g = child_module.m[0].cv2.conv.groups
    #         c2f_v2 = C3k2_infer(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=g,
    #                         c3k=c3k,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_v2_to_c2f(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     elif isinstance(child_module, C2PSA_v2):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f_v2 = C2PSA_infer(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), 
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_v2_to_c2f(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     else:
    #         replace_c2f_with_c2f_v2(child_module)

def infer_shortcut(bottleneck):
    try:
        c1 = bottleneck.cv1.conv.in_channels
        c2 = bottleneck.cv2.conv.out_channels
        return c1 == c2 and hasattr(bottleneck, 'add') and bottleneck.add
    except:
        return False

def transfer_weights_c2f_to_c2f_v2(c2f, c2f_v2):
    c2f_v2.cv2 = c2f.cv2
    c2f_v2.m = c2f.m

    state_dict = c2f.state_dict()
    state_dict_v2 = c2f_v2.state_dict()

    # Transfer cv1 weights from C2f to cv0 and cv1 in C2f_v2
    old_weight = state_dict['cv1.conv.weight']
    half_channels = old_weight.shape[0] // 2
    state_dict_v2['cv0.conv.weight'] = old_weight[:half_channels]
    state_dict_v2['cv1.conv.weight'] = old_weight[half_channels:]

    # Transfer cv1 batchnorm weights and buffers from C2f to cv0 and cv1 in C2f_v2
    for bn_key in ['weight', 'bias', 'running_mean', 'running_var']:
        old_bn = state_dict[f'cv1.bn.{bn_key}']
        state_dict_v2[f'cv0.bn.{bn_key}'] = old_bn[:half_channels]
        state_dict_v2[f'cv1.bn.{bn_key}'] = old_bn[half_channels:]

    # Transfer remaining weights and buffers
    for key in state_dict:
        if not key.startswith('cv1.'):
            state_dict_v2[key] = state_dict[key]

    # Transfer all non-method attributes
    for attr_name in dir(c2f):
        attr_value = getattr(c2f, attr_name)
        if not callable(attr_value) and '_' not in attr_name:
            setattr(c2f_v2, attr_name, attr_value)

    c2f_v2.load_state_dict(state_dict_v2)

def transfer_weights_elan_to_elan_v2(c2f, c2f_v2):
    c2f_v2.cv2 = c2f.cv2
    c2f_v2.cv3 = c2f.cv3
    c2f_v2.cv4 = c2f.cv4

    state_dict = c2f.state_dict()
    state_dict_v2 = c2f_v2.state_dict()

    # Transfer cv1 weights from C2f to cv0 and cv1 in C2f_v2
    old_weight = state_dict['cv1.conv.weight']
    half_channels = old_weight.shape[0] // 2
    state_dict_v2['cv0.conv.weight'] = old_weight[:half_channels]
    state_dict_v2['cv1.conv.weight'] = old_weight[half_channels:]

    # Transfer cv1 batchnorm weights and buffers from C2f to cv0 and cv1 in C2f_v2
    for bn_key in ['weight', 'bias', 'running_mean', 'running_var']:
        old_bn = state_dict[f'cv1.bn.{bn_key}']
        state_dict_v2[f'cv0.bn.{bn_key}'] = old_bn[:half_channels]
        state_dict_v2[f'cv1.bn.{bn_key}'] = old_bn[half_channels:]

    # Transfer remaining weights and buffers
    for key in state_dict:
        if not key.startswith('cv1.'):
            state_dict_v2[key] = state_dict[key]

    # Transfer all non-method attributes
    for attr_name in dir(c2f):
        attr_value = getattr(c2f, attr_name)
        if not callable(attr_value) and '_' not in attr_name:
            setattr(c2f_v2, attr_name, attr_value)

    c2f_v2.load_state_dict(state_dict_v2)

def replace_c2f_with_c2f_v2(module):
    # for yolov8n.yaml
    for name, child_module in module.named_children():
        if isinstance(child_module, C2f):
            # Replace C2f with C2f_v2 while preserving its parameters
            shortcut = infer_shortcut(child_module.m[0])
            c2f_v2 = C2f_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
                            n=len(child_module.m), shortcut=shortcut,
                            g=child_module.m[0].cv2.conv.groups,
                            e=child_module.c / child_module.cv2.conv.out_channels)
            transfer_weights_c2f_to_c2f_v2(child_module, c2f_v2)
            setattr(module, name, c2f_v2)
        else:
            replace_c2f_with_c2f_v2(child_module)
    
    # for yolov8-Faster-GFPN-P2-EfficientHead.yaml
    # for name, child_module in module.named_children():
    #     if isinstance(child_module, C2f_Faster):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f_v2 = C2f_Faster_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=1,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_to_c2f_v2(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     elif isinstance(child_module, C2f):
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f_v2 = C2f_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=child_module.m[0].cv2.conv.groups,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_to_c2f_v2(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     else:
    #         replace_c2f_with_c2f_v2(child_module)
    
    # for yolov8-BIFPN-EfficientRepHead.yaml
    # for name, child_module in module.named_children():
    #     if isinstance(child_module, C2f_EMBC):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f_v2 = C2f_EMBC_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=1,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_to_c2f_v2(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     elif isinstance(child_module, C2f):
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f_v2 = C2f_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=child_module.m[0].cv2.conv.groups,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_to_c2f_v2(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     else:
    #         replace_c2f_with_c2f_v2(child_module)
    
    # for yolov8-repvit-RepNCSPELAN.yaml
    # for name, child_module in module.named_children():
    #     if isinstance(child_module, RepNCSPELAN4):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         c2f_v2 = RepNCSPELAN4_v2(child_module.cv1.conv.in_channels, child_module.cv4.conv.out_channels,
    #                         child_module.cv1.conv.out_channels, child_module.cv3[-1].conv.out_channels, 1)
    #         transfer_weights_elan_to_elan_v2(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     else:
    #         replace_c2f_with_c2f_v2(child_module)
    
    # for yolov8-starnet-C2f-Star-LSCD.yaml
    # for name, child_module in module.named_children():
    #     if isinstance(child_module, C2f_Star):
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f_v2 = C2f_Star_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=1,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_to_c2f_v2(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     else:
    #         replace_c2f_with_c2f_v2(child_module)
    
    # for yolov10n.yaml
    # for name, child_module in module.named_children():
    #     if isinstance(child_module, C2fCIB):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f_v2 = C2fCIB_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         lk=(False if type(child_module.m[0].cv1[2]) is Conv else True),
    #                         g=1,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_to_c2f_v2(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     elif isinstance(child_module, C2f):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f_v2 = C2f_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=child_module.m[0].cv2.conv.groups,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_to_c2f_v2(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     else:
    #         replace_c2f_with_c2f_v2(child_module)
    
    # for yolo11n.yaml
    # for name, child_module in module.named_children():
    #     if isinstance(child_module, C3k2):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c3k = child_module.m[0] is C3k
    #         if c3k:
    #             g = child_module.m[0].m[0].cv2.conv.groups
    #         else:
    #             g = child_module.m[0].cv2.conv.groups
    #         c2f_v2 = C3k2_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), shortcut=shortcut,
    #                         g=g,
    #                         c3k=c3k,
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_to_c2f_v2(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     elif isinstance(child_module, C2PSA):
    #         # Replace C2f with C2f_v2 while preserving its parameters
    #         shortcut = infer_shortcut(child_module.m[0])
    #         c2f_v2 = C2PSA_v2(child_module.cv1.conv.in_channels, child_module.cv2.conv.out_channels,
    #                         n=len(child_module.m), 
    #                         e=child_module.c / child_module.cv2.conv.out_channels)
    #         transfer_weights_c2f_to_c2f_v2(child_module, c2f_v2)
    #         setattr(module, name, c2f_v2)
    #     else:
    #         replace_c2f_with_c2f_v2(child_module)


def transfer_weights_c2f_dynamic_to_c2f_dynamic_v2(old_module, new_module):
    """
    将 C2f_DynamicAxialRouterV2 的权重转换为 C2f_DynamicAxialRouterV2_v2 的权重。
    
    主要步骤：
      1. 直接转移 cv2、m 和 attention 模块（注意 m 内部结构应兼容）。
      2. 将 cv1 中的卷积权重按通道数一分为二，分别赋给新模块的 cv0 与 cv1。
      3. 对 cv1 中 BatchNorm 的参数同理进行切分。
      4. 转移除 cv1 外的其他参数（如果在目标 state_dict 中存在）。
      5. 仅转移非方法且属性名中不含下划线的公共属性，以避免影响后续剪枝。
    """
    # 直接转移 cv2、m 和 attention
    new_module.cv2 = old_module.cv2
    new_module.m = old_module.m
    new_module.attention = old_module.attention

    state_dict_old = old_module.state_dict()
    state_dict_new = new_module.state_dict()

    # 将 cv1 的卷积权重按通道数均分
    old_weight = state_dict_old['cv1.conv.weight']
    half_channels = old_weight.shape[0] // 2
    state_dict_new['cv0.conv.weight'] = old_weight[:half_channels].clone()
    state_dict_new['cv1.conv.weight'] = old_weight[half_channels:].clone()

    # 同步 BatchNorm 参数（weight, bias, running_mean, running_var）
    for bn_key in ['weight', 'bias', 'running_mean', 'running_var']:
        old_bn = state_dict_old[f'cv1.bn.{bn_key}']
        state_dict_new[f'cv0.bn.{bn_key}'] = old_bn[:half_channels].clone()
        state_dict_new[f'cv1.bn.{bn_key}'] = old_bn[half_channels:].clone()

    # 转移其他不以 'cv1.' 开头的参数（确保目标 state_dict 中存在）
    for key in state_dict_old:
        if not key.startswith('cv1.') and key in state_dict_new:
            state_dict_new[key] = state_dict_old[key]
    
    # 转移其他公共属性：仅转移那些非方法且名称中不含下划线的属性
    for attr_name in dir(old_module):
        try:
            attr_value = getattr(old_module, attr_name)
        except Exception:
            continue
        if not callable(attr_value) and ('_' not in attr_name):
            setattr(new_module, attr_name, attr_value)
    
    new_module.load_state_dict(state_dict_new)
    return new_module

# ----------------------------
# 递归替换模块中 C2f_DynamicAxialRouterV2 为 C2f_DynamicAxialRouterV2_v2
def replace_c2f_dynamic_with_c2f_dynamic_v2(module):
    """
    遍历 module 的所有子模块，将其中的 C2f_DynamicAxialRouterV2 实例替换为 C2f_DynamicAxialRouterV2_v2，
    并通过转换函数转移权重。
    """
    for name, child in module.named_children():
        if isinstance(child, C2f_DynamicAxialRouterV2):
            # 根据原模块的信息，计算 shortcut 和 e 参数
            shortcut = infer_shortcut(child)
            # e 通常为 child.c / child.cv2.conv.out_channels
            e = child.c / child.cv2.conv.out_channels
            # 尝试获取 m 中的动态路由参数，若无则使用默认值
            div_weight = child.m[0].div_weight if hasattr(child.m[0], 'div_weight') else 0.1
            temp_decay = child.m[0].temp_decay if hasattr(child.m[0], 'temp_decay') else 0.99
            new_instance = C2f_DynamicAxialRouterV2_v2(
                child.cv1.conv.in_channels, child.cv2.conv.out_channels,
                n=len(child.m),
                shortcut=shortcut,
                g=child.cv1.conv.groups,
                e=e,
                deploy=False,
                div_weight=div_weight,
                temp_decay=temp_decay
            )
            transfer_weights_c2f_dynamic_to_c2f_dynamic_v2(child, new_instance)
            setattr(module, name, new_instance)
        else:
            replace_c2f_dynamic_with_c2f_dynamic_v2(child)